#include <cstdio>
#include <stdlib.h>
#include <cuda.h>
#include <cuda_runtime.h>
#include <math.h>

#define A(i, j) A[i * Y + j]
#define B(i, j) B[i * Y + j]

void random_matrix(float *A, const int X, const int Y)
{
    for (int i = 0; i < X; i++)
        for (size_t j = 0; j < Y; j++)
            A(i, j) = 2.0 * (float)drand48() - 1.0;
}

void cpu_sgemm(float *A, float *B, float *C, const int M, const int K, const int N)
{
    for (size_t i = 0; i < M; i++)
    {
        for (size_t j = 0; j < N; j++)
        {
            float res = 0.f;
            for (size_t k = 0; k < K; k++)
            {
                res += A[i * K + k] * B[N * k + j];
            }
            C[i * N + j] = res;
        }
    }
}

float cal_diff(float *A, float *B, int X, int Y)
{
    int i, j;
    float max_diff = 0.0, diff;
    int printed = 0;
    for (i = 0; i < X; i++)
    {
        for (j = 0; j < Y; j++)
        {
            diff = abs(A(i, j) - B(i, j));
            max_diff = (diff > max_diff ? diff : max_diff);
            if (printed == 0)
            {
                if (max_diff > 0.5f)
                {
                    printf("error: i %d j %d diff %f got %f expect %f\n", i, j, max_diff, A(i, j), B(i, j));
                    printed = 1;
                }
            }
        }
    }
    return max_diff;
}

__global__ void cuda_sgemm(float *A, float *B, float *C, const int M, const int K, const int N)
{
    const int x = blockIdx.x * blockDim.x + threadIdx.x;
    const int y = blockIdx.y * blockDim.y + threadIdx.y;
    float *A_begin = A + K * blockIdx.y * blockDim.y;
    float *B_begin = B + blockIdx.x * blockDim.x;
    float res = 0.f;
    for (size_t k = 0; k < K; k++)
    {
        res += A_begin[threadIdx.y * K + k] * B_begin[k * N + threadIdx.x];
    }
    C[x + y * N] = res;
}

int main(int argc, char const *argv[])
{
    const int M = 512;
    const int N = 512;
    const int K = 512;

    size_t size_A = sizeof(float) * M * K;
    size_t size_B = sizeof(float) * N * K;
    size_t size_C = sizeof(float) * M * N;

    float *matrix_A_host = (float *)malloc(size_A);
    float *matrix_B_host = (float *)malloc(size_B);
    float *matrix_C_host_cpu = (float *)malloc(size_C);
    float *matrix_C_host_gpu = (float *)malloc(size_C);

    random_matrix(matrix_A_host, M, K);
    random_matrix(matrix_B_host, K, N);
    memset(matrix_C_host_cpu, 0.f, size_C);
    memset(matrix_C_host_gpu, 0.f, size_C);

    float *matrix_A_device;
    float *matrix_B_device;
    float *matrix_C_device;

    cudaMalloc((void **)&matrix_A_device, size_A);
    cudaMalloc((void **)&matrix_B_device, size_B);
    cudaMalloc((void **)&matrix_C_device, size_C);

    cudaMemcpy(matrix_A_device, matrix_A_host, size_A, cudaMemcpyHostToDevice);
    cudaMemcpy(matrix_B_device, matrix_B_host, size_B, cudaMemcpyHostToDevice);

    cpu_sgemm(matrix_A_host, matrix_B_host, matrix_C_host_cpu, M, K, N);

    constexpr int BLOCK = 16;
    dim3 block(BLOCK, BLOCK);
    dim3 grid((M + BLOCK - 1) / BLOCK, (N + BLOCK - 1) / BLOCK);

    cuda_sgemm<<<grid, block>>>(matrix_A_device, matrix_B_device, matrix_C_device, M, K, N);
    cudaMemcpy(matrix_C_host_gpu, matrix_C_device, size_C, cudaMemcpyDeviceToHost);

    float diff = cal_diff(matrix_C_host_gpu, matrix_C_host_cpu, M, N);

    if (diff > 0.5f)
    {
        printf("ans is wrong\n");
        exit(-1);
    }
    else
    {
        printf("ans is right\n");
    }

    free(matrix_A_host);
    free(matrix_B_host);
    free(matrix_C_host_cpu);
    free(matrix_C_host_gpu);

    cudaFree(matrix_A_device);
    cudaFree(matrix_B_device);
    cudaFree(matrix_C_device);

    return 0;
}
